import os, sys
import socket
import time

from paramiko import AuthenticationException, SFTPClient
from paramiko.client import SSHClient, AutoAddPolicy
from paramiko.ssh_exception import NoValidConnectionsError


class SshClient:
    @staticmethod
    def check_ssh_port(ip, port=22, username="pi", password="raspberry"):
        ssh_client = SSHClient()
        ssh_client.set_missing_host_key_policy(AutoAddPolicy())
        ssh_client.connect(host_ip, port=22, username=username, password=password)
        ssh_client.close()

    def __init__(self) -> None:
        self.ssh_client = SSHClient()

    def ssh_login(self, host_ip, username, password, timeout=5):
        try:
            self.ssh_client.set_missing_host_key_policy(AutoAddPolicy())
            self.ssh_client.connect(host_ip, port=22, username=username, password=password, timeout=timeout)
            self.shell = self.ssh_client.invoke_shell()
            self.sftp = self.ssh_client.open_sftp()
        except AuthenticationException:
            raise Exception("username or password error")
        except NoValidConnectionsError:
            raise Exception("connect time out")
        except:
            raise Exception(f"Unexpected error: {sys.exc_info()[0]}")

    def execute_some_command(self, command):
        stdin, stdout, stderr = self.ssh_client.exec_command(command)
        return stdout.read().decode()

    def execute_shell_command(self, command, wait_time=0.01):
        self.shell.send(command)
        time.sleep(wait_time)
        output = self.shell.recv(10000).decode("utf-8")
        return output

    # show upload progress
    def progress_bar(self, transferred, local_file_size, suffix=""):
        bar_len = 100
        filled_len = int(round(bar_len * transferred / float(local_file_size)))
        percents = round(100.0 * transferred / float(local_file_size), 1)
        bar_filled = "\033[32;1m%s\033[0m" % "*" * filled_len
        bar = bar_filled + "-" * (bar_len - filled_len)
        sys.stdout.write("[%s] %s%s %s\r" % (bar, "\033[32;1m%s\033[0m" % percents, "%", suffix))
        sys.stdout.flush()

    def __upload_file(self, local_file, remote_file, progress_bar):
        if not os.path.isfile(local_file):
            return

        # get file size
        local_file_size = os.stat(local_file).st_size
        print(f"local file:{local_file}({local_file_size}) copy to {remote_file}")
        with open(local_file, "rb") as f:
            # if progress_bar is None:
            #     progress_bar = self.progress_bar
            self.sftp.putfo(f, remote_file, local_file_size, callback=progress_bar)
            # sys.stdout.write("\n")

    def __path_type(self, path):
        result = self.execute_some_command(f"file {path}")
        if "cannot open" in result:
            return "None"
        elif ": directory" in result:
            return "directory"
        else:
            return "file"

    def upload_file_with_progress(self, local_path, remote_path, progress_bar=None, recursive=False):
        # ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        if os.path.isdir(local_path):
            for root, dirs, files in os.walk(local_path):
                relative_path = os.path.relpath(root, local_path)
                absolute_path = os.path.join(remote_path, relative_path)
                absolute_path = absolute_path.replace("\\", "/")
                self.execute_some_command(f"mkdir -p {absolute_path}")
                for file in files:
                    local_file = os.path.join(root, file)
                    remote_file = os.path.join(remote_path, os.path.relpath(local_file, local_path))
                    remote_file = remote_file.replace("\\", "/")
                    # print(remote_file)
                    self.__upload_file(local_file, remote_file, progress_bar)
                if not recursive:
                    for directory in dirs:
                        absolute_path = os.path.join(remote_path, directory)
                        absolute_path = absolute_path.replace("\\", "/")
                        self.execute_some_command(f"mkdir -p {absolute_path}")
                    break
        else:
            local_file = local_path
            result = self.__path_type(remote_path)
            if result == "directory":
                remote_file = os.path.join(remote_path, os.path.basename(local_file))
            else:
                remote_file = remote_path
            remote_file = remote_file.replace("\\", "/")
            self.__upload_file(local_file, remote_file, progress_bar)

    def __download_file(self, remote_file, local_file, progress_bar):
        with open(local_file, "wb") as file:
            self.sftp.getfo(remote_file, file, callback=progress_bar)

    def download_file_with_progress(self, remote_path, local_path, progress_bar=None):
        # self.__download_file(remote_path, local_path, progress_bar)

        path_type = self.__path_type(remote_path)
        if path_type == "directory":
            # 创建目录
            remote_dirs = self.execute_some_command(f"find {remote_path} -type d")
            relative_dirs = remote_dirs.replace(remote_path, "").splitlines()
            for relative_dir in relative_dirs:
                local_dir = os.path.join(local_path, relative_dir[1:]).replace("\\", "/")
                if not os.path.exists(local_dir):
                    # print(local_dir)
                    os.makedirs(local_dir)

            # 复制文件
            remote_files = self.execute_some_command(f"find {remote_path} -type f")
            relative_files = remote_files.replace(remote_path, "").splitlines()
            for relative_file in relative_files:
                remote_file = os.path.join(remote_path, relative_file[1:]).replace("\\", "/")
                local_file = os.path.join(local_path, relative_file[1:]).replace("\\", "/")
                # print(remote_file, local_file)
                self.__download_file(remote_file, local_file, progress_bar)

        else:
            # 复制文件
            remote_file = remote_path
            if os.path.isdir(local_path):
                basename = os.path.basename(remote_file)
                local_file = os.path.join(local_path, basename)
                local_file = local_file.replace("\\", "/")
            else:
                local_file = local_path

            self.__download_file(remote_file, local_file, progress_bar)

    def ssh_logout(self):
        self.sftp.close()
        self.shell.close()
        self.ssh_client.close()


if __name__ == "__main__":
    temp_dir = f"__FS620_temp"
    compressed_file = f"FS620.tar.gz"
    app_dir = "/home/pi/app"
    app_name_head = "FS620_"

    ip_address = "192.168.3.112"

    ssh = SshClient()
    ssh.ssh_login(ip_address, "pi", "raspberry")

    ssh.download_file_with_progress("/home/pi/temp/update_list2.json", "temp/update_list2.json", progress_bar=ssh.progress_bar)
    ssh.download_file_with_progress("/home/pi/temp/update_list2.json", "temp/", progress_bar=ssh.progress_bar)
    # ssh.download_file_with_progress("/home/pi/temp/1", "temp", progress_bar=ssh.progress_bar)

    # ssh.upload_file_with_progress(f"download/xxxx", "/home/pi/temp")  # 文件不存在
    # ssh.upload_file_with_progress(f"download/update_list.json", "/home/pi/temp/update_list2.json")  # 单文件2单文件
    # ssh.upload_file_with_progress(f"download/update_list.json", "/home/pi/temp")  # 单文件2目录
    # ssh.upload_file_with_progress(f"d:/iolmaster/update_app/temp", "/home/pi/temp/1", recursive=False, progress_bar=ssh.progress_bar)  # 不递归
    # ssh.upload_file_with_progress(f"d:/iolmaster/update_app/temp", "/home/pi/temp/1", recursive=True, progress_bar=ssh.progress_bar)  # 递归

    ssh.ssh_logout()

    # # 远程启动应用程序 调试
    # current = "V1.0.26"
    # subprocess.getoutput(f"scp -o StrictHostKeyChecking=no app/startup_app.sh pi@{ip_address}:{app_dir}")
    # ssh.execute_shell_command(f"cd {app_dir}\n")
    # ssh.execute_shell_command(f"chmod +x ./startup_app.sh\n")
    # print(ssh.execute_shell_command(f"ls\n",1))
